"""
model name : 深度学习
file       : dataset.py
information:
    author : OuYang
    time   : 2025/1/13
"""

import os

import torch
import torchvision
from PIL import Image
import numpy as np
from torch.utils.data import Dataset


# Yolo Dataset
class YOLODataset(Dataset):
    def __init__(self,
                 root,
                 num_classes=3,
                 s=7,
                 b=2,
                 train=True,
                 transform=None
                 ):
        super(YOLODataset, self).__init__()
        self.root = root
        self.num_classes = num_classes
        self.s = s
        self.b = b
        self.train = train
        self.transform = transform

        # Cal relevant arguments
        self.cell_size = 1.0 / self.s

        # Train Data or Valid Data
        if self.train:
            self.image_dir = os.path.join(self.root, 'train/images')
            self.label_dir = os.path.join(self.root, 'train/labels')
        else:
            self.image_dir = os.path.join(self.root, 'valid/images')
            self.label_dir = os.path.join(self.root, 'valid/labels')

        # Load Data
        self.images = os.listdir(self.image_dir)

    def __len__(self):
        return len(self.images)

    def __getitem__(self, index):
        image_path = os.path.join(self.image_dir, self.images[index])
        label_path = os.path.join(self.label_dir, self.images[index].split('.')[0] + '.txt')

        # Open Image
        image = Image.open(image_path).convert('RGB')

        # transform
        if self.transform:
            image = self.transform(image)
        else:
            image = torchvision.transforms.ToTensor()(image)

        # Load Label
        labels = []
        if os.path.exists(label_path):
            with open(label_path, 'r') as file:
                for line in file:
                    class_index, x, y, w, h = line.strip().strip().split(' ')
                    labels.append([
                        int(class_index),
                        float(x),
                        float(y),
                        float(w),
                        float(h)
                    ])

        target = self._build_targets(labels)

        return image, target

    def _build_targets(self, labels):
        target = torch.zeros((self.s, self.s, self._cal_one_grid_predictions()))
        for label in labels:
            class_index, x, y, w, h = label

            i, j = int(x // self.cell_size), int(y // self.cell_size)
            x = x / self.cell_size - i
            y = y / self.cell_size - j

            for bi in range(self.b):
                target[i][j][bi * 5:bi * 5 + 4] = torch.tensor([x, y, w, h])
                target[i][j][bi * 5 + 4] = 1
            target[i][j][self.b * 5 + class_index] = 1

        return target

    def _cal_one_grid_predictions(self):
        return self.b * 5 + self.num_classes


if __name__ == '__main__':
    from torchvision import transforms

    dataset = YOLODataset(
        root='d:/space/datasets/TeaProjectDatasetVersion/Test',
        num_classes=1,
        transform=transforms.Compose([
            transforms.Resize(448),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        ])
    )

    data = dataset[0]

    image, target = data
    print(target)
    import cv2
    image = np.array(image)
    image = np.transpose(image, (1, 2, 0))
    image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)

    cell_size = 1.0 / 7
    print(cell_size)
    target_x = 0.6203
    target_y = 0.3688
    target_w = 0.4500
    target_h = 0.3219
    target_x1 = target_x * cell_size + cell_size * 3 - target_w / 2
    target_y1 = target_y * cell_size + cell_size * 3 - target_h / 2
    target_x2 = target_x * cell_size + cell_size * 3 + target_w / 2
    target_y2 = target_y * cell_size + cell_size * 3 + target_h / 2

    cv2.rectangle(
        image,
        (int(target_x1 * 448), int(target_y1 * 448)),
        (int(target_x2 * 448), int(target_y2 * 448)),
        (0, 0, 255)  # 红色目标框
    )

    cv2.imshow('image', image)
    cv2.waitKey(0)
    cv2.destroyAllWindows()

    print(target.shape)
    print(target)
